package com.example.util.tree;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

public class TreeBuilder {

    /**
     * 构建树形结构
     *
     * @param list 数据列表
     * @param idGetter 节点ID获取函数
     * @param parentIdGetter 父节点ID获取函数
     * @param rootParentId 根节点的父ID值
     * @param <T> 数据类型
     * @param <K> ID类型
     * @return 树形结构列表
     */
    public static <T, K> List<TreeNode<T>> buildTree(List<T> list,
                                                     Function<T, K> idGetter,
                                                     Function<T, K> parentIdGetter,
                                                     K rootParentId) {
        if (list == null || list.isEmpty()) {
            return new ArrayList<>();
        }

        // 将列表转换为Map，key为节点ID，value为TreeNode
        Map<K, TreeNode<T>> nodeMap = list.stream()
                .collect(Collectors.toMap(idGetter, TreeNode::new));

        List<TreeNode<T>> roots = new ArrayList<>();

        // 构建树形结构
        for (T item : list) {
            K id = idGetter.apply(item);
            K parentId = parentIdGetter.apply(item);
            TreeNode<T> node = nodeMap.get(id);

            if (isRootNode(parentId, rootParentId)) {
                // 根节点
                roots.add(node);
            } else {
                // 子节点，找到父节点并添加
                TreeNode<T> parentNode = nodeMap.get(parentId);
                if (parentNode != null) {
                    parentNode.addChild(node);
                } else {
                    // 如果父节点不存在，也作为根节点处理
                    roots.add(node);
                }
            }
        }

        return roots;
    }

    /**
     * 判断是否为根节点
     */
    private static <K> boolean isRootNode(K parentId, K rootParentId) {
        if (rootParentId == null) {
            return parentId == null;
        }
        return rootParentId.equals(parentId);
    }

    /**
     * 构建树形结构（支持自定义根节点条件）
     *
     * @param list 数据列表
     * @param idGetter 节点ID获取函数
     * @param parentIdGetter 父节点ID获取函数
     * @param isRootPredicate 根节点判断条件
     * @param <T> 数据类型
     * @param <K> ID类型
     * @return 树形结构列表
     */
    public static <T, K> List<TreeNode<T>> buildTree(List<T> list,
                                                     Function<T, K> idGetter,
                                                     Function<T, K> parentIdGetter,
                                                     java.util.function.Predicate<T> isRootPredicate) {
        if (list == null || list.isEmpty()) {
            return new ArrayList<>();
        }

        Map<K, TreeNode<T>> nodeMap = list.stream()
                .collect(Collectors.toMap(idGetter, TreeNode::new));

        List<TreeNode<T>> roots = new ArrayList<>();

        for (T item : list) {
            K id = idGetter.apply(item);
            K parentId = parentIdGetter.apply(item);
            TreeNode<T> node = nodeMap.get(id);

            if (isRootPredicate.test(item)) {
                roots.add(node);
            } else {
                TreeNode<T> parentNode = nodeMap.get(parentId);
                if (parentNode != null) {
                    parentNode.addChild(node);
                } else {
                    roots.add(node);
                }
            }
        }

        return roots;
    }
}
